#include "darts_adjust_node/darts_adjust_node.hpp"
#include <tf2_ros/create_timer_ros.h>

namespace darts_adjust
{

// 在构造函数中添加发布者初始化
DartsAdjustNode::DartsAdjustNode(const rclcpp::NodeOptions & options)
: Node("darts_adjust_node", options)
{
  RCLCPP_INFO(this->get_logger(), "Hello DartsAdjustNode!");

  // 初始化相机内参矩阵
  camera_matrix_ = cv::Mat::zeros(3, 3, CV_64F);
  
  // 初始化帧率计数器
  frame_count_ = 0;
  last_time_ = this->now();
  
  // 初始化卡尔曼滤波器状态
  kalman_initialized_ = false;
  initKalmanFilter();
  
  // 创建误差发布者
  error_pub_ = this->create_publisher<std_msgs::msg::Float64>("/darts/error", 10);
  
  // 创建调试图像发布者
  debug_image_pub_ = this->create_publisher<sensor_msgs::msg::Image>("/darts/debug_image", 10);
  
  // 订阅相机内参
  camera_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
      "/camera_info", rclcpp::SensorDataQoS(),
      [this](sensor_msgs::msg::CameraInfo::ConstSharedPtr camera_info) {
          RCLCPP_INFO(this->get_logger(), "Received camera info!");
          RCLCPP_INFO(this->get_logger(), "K matrix: [%f, %f, %f, %f, %f, %f, %f, %f, %f]",
              camera_info->k[0], camera_info->k[1], camera_info->k[2],
              camera_info->k[3], camera_info->k[4], camera_info->k[5],
              camera_info->k[6], camera_info->k[7], camera_info->k[8]);
          camera_matrix_.at<double>(0, 0) = camera_info->k[0];  // fx
          camera_matrix_.at<double>(0, 2) = camera_info->k[2];   // cx
          camera_matrix_.at<double>(1, 1) = camera_info->k[4];   // fy
          camera_matrix_.at<double>(1, 2) = camera_info->k[5];    // cy
          camera_matrix_.at<double>(2, 2) = 1.0;
          camera_info_sub_.reset();
  });

  // 创建订阅者，订阅图像原始数据（/image_raw）
  darts_image_sub_ = create_subscription<sensor_msgs::msg::Image>(
    "/image_raw", rclcpp::SensorDataQoS(),
    std::bind(&DartsAdjustNode::imageCallback, this, std::placeholders::_1));
}

void DartsAdjustNode::imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr & msg)
{
  // 帧率计算
  auto current_time = this->now();
  frame_count_++;
  double elapsed = (current_time - last_time_).seconds();
  
  if (elapsed >= 1.0) {
      double fps = frame_count_ / elapsed;
      RCLCPP_INFO(this->get_logger(), "[FPS] Current: %.2f", fps);
      frame_count_ = 0;
      last_time_ = current_time;
  }

  cv::Mat darts_image;
  try
  {
      // 使用cv_bridge将ROS图像消息转换为OpenCV图像
      cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::RGB8);
      darts_image = cv_ptr->image;  // 修正变量名：rune_image -> darts_image
  }
  catch (cv_bridge::Exception & e)
  {
      // 如果图像转换失败，输出错误信息
      RCLCPP_ERROR(get_logger(), "cv_bridge exception: %s", e.what());

      // 重置卡尔曼滤波
      kalman_initialized_ = false;
      target_error_ = 0;
      return;
  }

  // 图像处理
  cv::Mat processed_image = processImage(darts_image);  // 修正变量名
  
  // 发布调试图像
  if (debug_image_pub_->get_subscription_count() > 0) {
    cv_bridge::CvImage debug_img;
    debug_img.header = msg->header;
    debug_img.encoding = sensor_msgs::image_encodings::RGB8;
    debug_img.image = processed_image;
    debug_image_pub_->publish(*debug_img.toImageMsg());
    RCLCPP_DEBUG(get_logger(), "Published debug image");
  }
  
  // 发布误差
  std_msgs::msg::Float64 error_msg;
  error_msg.data = static_cast<double>(target_error_);
  error_pub_->publish(error_msg);
}

cv::Mat DartsAdjustNode::processImage(const cv::Mat & input_image)
{
  // 创建输出图像
  cv::Mat output_image = input_image.clone();
  
  // 转换到HSV颜色空间进行颜色阈值处理
  cv::Mat hsv_image;
  cv::cvtColor(input_image, hsv_image, cv::COLOR_RGB2HSV);
  
  // 定义绿色阈值范围（调整后的参数）
  cv::Scalar lower_green(35, 100, 100);     // 绿色范围下界 (H:35~85)
  cv::Scalar upper_green(85, 255, 255);     // 绿色范围上界
  
  // 创建掩码
  cv::Mat mask;
  cv::inRange(hsv_image, lower_green, upper_green, mask); 
  
  // 应用掩码提取目标区域
  cv::Mat result;
  cv::bitwise_and(input_image, input_image, result, mask);
  
  // 形态学操作优化掩码
  cv::Mat kernel = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(5, 5));
  cv::morphologyEx(mask, mask, cv::MORPH_OPEN, kernel);  // 开运算去除噪点
  cv::morphologyEx(mask, mask, cv::MORPH_CLOSE, kernel); // 闭运算填充小洞
  
  // 寻找轮廓
  std::vector<std::vector<cv::Point>> contours;
  std::vector<cv::Vec4i> hierarchy;
  cv::findContours(mask, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
  
  // 绘制轮廓
  cv::drawContours(output_image, contours, -1, cv::Scalar(0, 255, 0), 2);
  
  // 如果找到轮廓，处理最大的轮廓（假设是目标）
  if (!contours.empty()) {
    // 找到最大轮廓
    size_t max_contour_idx = 0;
    double max_area = 0;
    for (size_t i = 0; i < contours.size(); i++) {
      double area = cv::contourArea(contours[i]);
      if (area > max_area) {
        max_area = area;
        max_contour_idx = i;
      }
    }
    
    // 计算最大轮廓的边界框
    cv::Rect bounding_rect = cv::boundingRect(contours[max_contour_idx]);
    cv::rectangle(output_image, bounding_rect, cv::Scalar(255, 0, 0), 2);
    
    // 计算并显示目标中心
    cv::Point center(bounding_rect.x + bounding_rect.width/2, 
                     bounding_rect.y + bounding_rect.height/2);
    
    // 计算目标中心与图像中心的x轴像素差
    int image_center_x = input_image.cols / 2;
    int raw_error = center.x - image_center_x;
    
    // 应用卡尔曼滤波
    measurement_.at<float>(0) = static_cast<float>(raw_error);
    
    if (!kalman_initialized_) {
      // 首次检测到目标，初始化卡尔曼滤波器状态
      kf_.statePost.at<float>(0) = static_cast<float>(raw_error);
      kf_.statePost.at<float>(1) = 0;  // 初始速度为0
      kalman_initialized_ = true;
    } else {
      // 预测
      cv::Mat prediction = kf_.predict();
      // 更新
      cv::Mat estimated = kf_.correct(measurement_);
      // 使用滤波后的结果
      target_error_ = static_cast<int>(estimated.at<float>(0));
      
      // 在图像上显示原始位置和滤波后的位置
      cv::Point filtered_center(image_center_x + target_error_, center.y);
      cv::circle(output_image, center, 5, cv::Scalar(0, 0, 255), -1);  // 原始位置（红色）
      cv::circle(output_image, filtered_center, 8, cv::Scalar(255, 0, 255), 2);  // 滤波后位置（紫色）
    }
    
    // 在图像上绘制图像中心线
    cv::line(output_image, cv::Point(image_center_x, 0), 
             cv::Point(image_center_x, input_image.rows), 
             cv::Scalar(255, 255, 0), 1);
    
    // 显示目标信息和误差值
    std::string info = "Target: (" + std::to_string(center.x) + ", " + 
                       std::to_string(center.y) + ")";
    std::string error_info = "Error X: " + std::to_string(target_error_) + " px";
    std::string raw_error_info = "Raw Error X: " + std::to_string(raw_error) + " px";
    cv::putText(output_image, info, cv::Point(10, 30), 
                cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 0), 2);
    cv::putText(output_image, error_info, cv::Point(10, 60), 
                cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 255, 0), 2);
    cv::putText(output_image, raw_error_info, cv::Point(10, 90), 
                cv::FONT_HERSHEY_SIMPLEX, 0.7, cv::Scalar(0, 0, 255), 2);
  }
  else {
    // 如果没有找到目标，将误差设为0
    target_error_ = 10000;
    kalman_initialized_ = false;  // 重置卡尔曼滤波器状态
  }
  
  // 添加处理后的掩码到输出图像的一角
  cv::Mat mask_small;
  cv::resize(mask, mask_small, cv::Size(input_image.cols/4, input_image.rows/4));
  cv::cvtColor(mask_small, mask_small, cv::COLOR_GRAY2RGB);
  cv::Mat roi = output_image(cv::Rect(0, 0, mask_small.cols, mask_small.rows));
  cv::addWeighted(roi, 0.7, mask_small, 0.3, 0, roi);
  
  // 添加文本
  cv::putText(output_image, "Darts Adjust", cv::Point(50, 70), 
              cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 2);
  
  return output_image;
}

void DartsAdjustNode::initKalmanFilter()
{
  // 状态变量：[x, vx]，其中x是位置，vx是速度
  kf_ = cv::KalmanFilter(2, 1, 0);
  
  // 状态转移矩阵 A
  // [1, dt]
  // [0, 1]
  kf_.transitionMatrix = (cv::Mat_<float>(2, 2) << 1, 1, 0, 1);
  
  // 测量矩阵 H
  // [1, 0] - 只测量位置，不测量速度
  kf_.measurementMatrix = (cv::Mat_<float>(1, 2) << 1, 0);
  
  // 过程噪声协方差矩阵 Q
  // 较小的值表示对模型的高信任度
  kf_.processNoiseCov = (cv::Mat_<float>(2, 2) << 1e-4, 0, 0, 1e-4);
  
  // 测量噪声协方差矩阵 R
  // 较大的值表示对测量的低信任度
  kf_.measurementNoiseCov = (cv::Mat_<float>(1, 1) << 1e-1);
  
  // 后验错误估计协方差矩阵 P
  kf_.errorCovPost = (cv::Mat_<float>(2, 2) << 1, 0, 0, 1);
  
  // 初始化状态
  kf_.statePost = (cv::Mat_<float>(2, 1) << 0, 0);
  
  // 初始化测量值
  measurement_ = cv::Mat(1, 1, CV_32F);
}

}  // namespace darts_adjust

// 主函数
int main(int argc, char * argv[])
{
  rclcpp::init(argc, argv);
  auto node = std::make_shared<darts_adjust::DartsAdjustNode>();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}